#include <cuda.h>
#include <cuda_runtime.h>
#include <curand_kernel.h>

#include <cstdio>
#include <stdint.h>
#include <stdexcept>
#include <limits>
#include <algorithm>

#include "rayutil.h"

template <typename T>
inline __host__ __device__ T div_round_up(T val, T divisor) {
	return (val + divisor - 1) / divisor;
}

inline __device__ void swapf(float& a, float& b) {
	float c = a; a = b; b = c;
}
inline __device__ float clamp(const float x, const float min, const float max) {
	return fminf(max, fmaxf(min, x));
}

template <typename scalar_t>
__global__ void kernel_rays_coords(
	const scalar_t* __restrict__ rays_o,
	const scalar_t* __restrict__ rays_d,
	const float* __restrict__ aabb,
	const uint32_t num_steps, const float min_near, const uint32_t N,
	const bool perturb,
	scalar_t* __restrict__ nears, 
	scalar_t* __restrict__ fars,
	scalar_t* __restrict__ xyzs,
	scalar_t* __restrict__ z_vals,
	scalar_t* __restrict__ deltas
)
{
	const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
	if (n >= N) return;

	rays_o += n * 3;
	rays_d += n * 3;
	xyzs += n * num_steps * 3;
	deltas += n * num_steps;
	z_vals += n * num_steps;

	curandStateXORWOW_t state{};
	curand_init(threadIdx.x, n, 0, &state);

	const float ox = rays_o[0], oy = rays_o[1], oz = rays_o[2];
	const float dx = rays_d[0], dy = rays_d[1], dz = rays_d[2];
	const float minx = aabb[0], miny = aabb[1], minz = aabb[2];
	const float maxx = aabb[3], maxy = aabb[4], maxz = aabb[5];
	const float rdx = 1 / dx, rdy = 1 / dy, rdz = 1 / dz;

	// get near far (assume cube scene)
	float near = (aabb[0] - ox) * rdx;
	float far = (aabb[3] - ox) * rdx;
	if (near > far) swapf(near, far);

	float near_y = (aabb[1] - oy) * rdy;
	float far_y = (aabb[4] - oy) * rdy;
	if (near_y > far_y) swapf(near_y, far_y);

	if (near > far_y || near_y > far) {
		nears[n] = fars[n] = std::numeric_limits<scalar_t>::max();
		return;
	}

	if (near_y > near) near = near_y;
	if (far_y < far) far = far_y;

	float near_z = (aabb[2] - oz) * rdz;
	float far_z = (aabb[5] - oz) * rdz;
	if (near_z > far_z) swapf(near_z, far_z);

	if (near > far_z || near_z > far) {
		nears[n] = fars[n] = std::numeric_limits<scalar_t>::max();
		return;
	}

	if (near_z > near) near = near_z;
	if (far_z < far) far = far_z;

	if (near < min_near) near = min_near;

	nears[n] = near;
	fars[n] = far;

	float step = 1.0 / num_steps;
	float sample_dist = (far - near) / num_steps;
	float dt = sample_dist;
	for (int t = 0; t < num_steps; t++)
	{
		float ct = t * step;
		float zv = near + (far - near) * ct;
		if (perturb)
		{
			float rand = curand_uniform(&state);
			zv = zv + (rand - 0.5) * sample_dist;
		}
		float x = ox + dx * zv;
		float y = oy + dy * zv;
		float z = oz + dz * zv;
		x = clamp(x, minx, maxx);
		y = clamp(y, miny, maxy);
		z = clamp(z, minz, maxz);
		if (t > 0) {
			dt = zv - z_vals[t - 1];
			deltas[t - 1] = dt;
		}
		if (t == num_steps - 1) {
			deltas[t] = sample_dist;
		}
		xyzs[t * 3] = x;
		xyzs[t * 3 + 1] = y;
		xyzs[t * 3 + 2] = z;
		z_vals[t] = zv;
	}
}


template <typename scalar_t>
void rays_coords_wrapper(const scalar_t* rays_o, const scalar_t* rays_d, const float* aabb,
	const uint32_t num_steps, const float min_near, const uint32_t N, const bool perturb,
	scalar_t* nears, scalar_t* fars, scalar_t* xyzs, scalar_t* z_vals, scalar_t* deltas)
{
	static constexpr uint32_t N_THREAD = 128;
	kernel_rays_coords<scalar_t> << <div_round_up(N, N_THREAD), N_THREAD >> > (rays_o, rays_d, aabb,
		num_steps, min_near, N, perturb, nears, fars, xyzs, z_vals, deltas);
}

void rays_coords(const at::Tensor* rays_o, const at::Tensor* rays_d, const at::Tensor* aabb,
	const uint32_t num_steps, const float min_near, const uint32_t N, const bool perturb,
	at::Tensor* nears, at::Tensor* fars, at::Tensor* xyzs, at::Tensor* z_vals, at::Tensor* deltas)
{
	AT_DISPATCH_FLOATING_TYPES_AND_HALF(
		rays_o->scalar_type(), "near_far_from_aabb", ([&] {
			rays_coords_wrapper(rays_o->data_ptr<scalar_t>(), rays_d->data_ptr<scalar_t>(), aabb->data_ptr<float>(),
				num_steps, min_near, N, perturb,
				nears->data_ptr<scalar_t>(), fars->data_ptr<scalar_t>(),
				xyzs->data_ptr<scalar_t>(), z_vals->data_ptr<scalar_t>(), deltas->data_ptr<scalar_t>());
			}));
}